package com.liru.neo4j_demo.util;

import com.liru.neo4j_demo.model.Chunk;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 文件内容切割工具类
 * 提供按行、段落、字符数等规则切割文件内容的功能
 */
public class FileSplitter {

    // 假设1个中文字符约等于1个token
    private static final int TOKEN_THRESHOLD = 500;
    private static final int MIN_CHUNK_SIZE = 200;
    private static final int MAX_CHUNK_SIZE = 400;
    private static final double SIMILARITY_THRESHOLD = 0.5;

    /**
     * 按行数切割文本
     * @param text 待切割的文本
     * @param linesPerSegment 每段包含的行数
     * @return 切割后的文本片段列表
     */
    public static List<String> splitByLines(String text, int linesPerSegment) {
        List<String> segments = new ArrayList<>();
        String[] lines = text.split("\\n");

        StringBuilder segment = new StringBuilder();
        for (int i = 0; i < lines.length; i++) {
            segment.append(lines[i]).append("\n");

            if ((i + 1) % linesPerSegment == 0 || i == lines.length - 1) {
                segments.add(segment.toString().trim());
                segment.setLength(0);
            }
        }

        return segments;
    }

    /**
     * 按段落切割文本
     * @param text 待切割的文本
     * @return 切割后的文本片段列表
     */
    public static List<String> splitByParagraphs(String text) {
        List<String> segments = new ArrayList<>();
        // 使用正则表达式匹配段落（两个或多个换行符分隔）
        String[] paragraphs = text.split("\\n+");

        for (String paragraph : paragraphs) {
            if (!paragraph.trim().isEmpty()) {
                segments.add(paragraph.trim());
            }
        }

        return segments;
    }

    /**
     * 按字符数切割文本
     * @param text 待切割的文本
     * @param charsPerSegment 每段包含的字符数
     * @return 切割后的文本片段列表
     */
    public static List<String> splitByCharacters(String text, int charsPerSegment) {
        List<String> segments = new ArrayList<>();
        int textLength = text.length();

        for (int i = 0; i < textLength; i += charsPerSegment) {
            int end = Math.min(i + charsPerSegment, textLength);
            segments.add(text.substring(i, end));
        }

        return segments;
    }

    /**
     * 按自定义正则表达式切割文本
     * @param text 待切割的文本
     * @param regex 切割正则表达式
     * @return 切割后的文本片段列表
     */
    public static List<String> splitByRegex(String text, String regex) {
        List<String> segments = new ArrayList<>();
        Pattern pattern = Pattern.compile(regex);
        String[] parts = pattern.split(text);

        for (String part : parts) {
            if (!part.trim().isEmpty()) {
                segments.add(part.trim());
            }
        }

        return segments;
    }

    /**
     * 高级文本切割策略
     * 1. 粗切：先按段落（\n\n）切
     * 2. 如果段落 > 500 token（约 300~400 字），则进入细切步骤
     * 3. 细切：按句子标点（。？！）切成小句
     * 4. 再用滑动窗口组合成 200~400 token 的块
     * 5. 可选语义优化：对相邻句子做向量相似度计算（如相似度 < 0.5 则断开）
     *
     * @param text 待切割文本
     * @param embeddingModel 嵌入模型（用于语义优化）
     * @return 切割后的文本片段列表
     */
    public static List<String> splitByAdvancedStrategy(String text, EmbeddingModel embeddingModel) {
        List<String> result = new ArrayList<>();

        // 1. 粗切：按段落（\n\n）切割
        String[] paragraphs = text.split("\\n+");

        for (String paragraph : paragraphs) {
            paragraph = paragraph.trim();
            if (paragraph.isEmpty()) continue;

            // 估算段落的token数（假设1个中文字符约等于1个token）
            int tokenCount = paragraph.length();

            // 2. 如果段落 <= 500 token，直接添加到结果
            if (tokenCount <= TOKEN_THRESHOLD) {
                result.add(paragraph);
            } else {
                // 3. 细切：按句子标点（。？！）切成小句
                List<String> sentences = splitSentences(paragraph);

                // 4. 用滑动窗口组合成 200~400 token 的块
                List<String> chunks = createSlidingWindowChunks(sentences);

                // 5. 可选语义优化
                if (embeddingModel != null) {
                    chunks = optimizeWithSemanticSimilarity(chunks, embeddingModel);
                }

                result.addAll(chunks);
            }
        }

        return result;
    }

    /**
     * 按句子标点（。？！）切割文本
     * @param text 待切割文本
     * @return 句子列表
     */
    private static List<String> splitSentences(String text) {
        List<String> sentences = new ArrayList<>();
        // 使用正则表达式匹配中文句子标点
        Pattern pattern = Pattern.compile("[^。？！]*[。？！]");
        Matcher matcher = pattern.matcher(text);

        while (matcher.find()) {
            sentences.add(matcher.group().trim());
        }

        // 处理没有结束标点的最后一句
        int lastEnd = 0;
        while (lastEnd < text.length() && matcher.find(lastEnd)) {
            lastEnd = matcher.end();
        }
        if (lastEnd < text.length()) {
            String lastSentence = text.substring(lastEnd).trim();
            if (!lastSentence.isEmpty()) {
                sentences.add(lastSentence);
            }
        }

        return sentences;
    }

    /**
     * 用滑动窗口将句子组合成块
     * @param sentences 句子列表
     * @return 组合后的块列表
     */
    private static List<String> createSlidingWindowChunks(List<String> sentences) {
        List<String> chunks = new ArrayList<>();
        int start = 0;

        while (start < sentences.size()) {
            StringBuilder chunkBuilder = new StringBuilder();
            int currentSize = 0;
            int end = start;

            // 添加句子直到达到最大token数
            while (end < sentences.size()) {
                String sentence = sentences.get(end);
                int sentenceSize = sentence.length();

                // 如果添加当前句子会超过最大token数，则停止
                if (currentSize + sentenceSize > MAX_CHUNK_SIZE) {
                    break;
                }

                chunkBuilder.append(sentence);
                currentSize += sentenceSize;
                end++;
            }

            // 如果当前块太小，尝试添加更多句子
            if (currentSize < MIN_CHUNK_SIZE && end < sentences.size()) {
                while (end < sentences.size()) {
                    String sentence = sentences.get(end);
                    chunkBuilder.append(sentence);
                    currentSize += sentence.length();
                    end++;

                    if (currentSize >= MIN_CHUNK_SIZE) {
                        break;
                    }
                }
            }

            chunks.add(chunkBuilder.toString().trim());
            start = end;
        }

        return chunks;
    }

    /**
     * 使用语义相似度优化块分割
     * @param chunks 块列表
     * @param embeddingModel 嵌入模型
     * @return 优化后的块列表
     */
    private static List<String> optimizeWithSemanticSimilarity(List<String> chunks, EmbeddingModel embeddingModel) {
        List<String> optimizedChunks = new ArrayList<>();

        if (chunks.isEmpty()) {
            return optimizedChunks;
        }

        optimizedChunks.add(chunks.get(0));

        for (int i = 1; i < chunks.size(); i++) {
            String previousChunk = optimizedChunks.get(optimizedChunks.size() - 1);
            String currentChunk = chunks.get(i);

            // 计算相邻块的语义相似度
            double similarity = calculateSimilarity(previousChunk, currentChunk, embeddingModel);

            // 如果相似度低于阈值，则将当前块作为新块
            if (similarity < SIMILARITY_THRESHOLD) {
                optimizedChunks.add(currentChunk);
            } else {
                // 否则，合并到前一个块
                String mergedChunk = previousChunk + " " + currentChunk;
                optimizedChunks.set(optimizedChunks.size() - 1, mergedChunk);
            }
        }

        return optimizedChunks;
    }

    /**
     * 计算两个文本的语义相似度
     * @param text1 第一个文本
     * @param text2 第二个文本
     * @param embeddingModel 嵌入模型
     * @return 相似度分数
     */
    private static double calculateSimilarity(String text1, String text2, EmbeddingModel embeddingModel) {
        Embedding embedding1 = embeddingModel.embed(text1).content();
        Embedding embedding2 = embeddingModel.embed(text2).content();

        // 计算余弦相似度
        List<Float> vec1 = embedding1.vectorAsList();
        List<Float> vec2 = embedding2.vectorAsList();

        double dotProduct = 0.0;
        double norm1 = 0.0;
        double norm2 = 0.0;

        for (int i = 0; i < vec1.size() && i < vec2.size(); i++) {
            dotProduct += vec1.get(i) * vec2.get(i);
            norm1 += Math.pow(vec1.get(i), 2);
            norm2 += Math.pow(vec2.get(i), 2);
        }

        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }
}